package com.example.myspringai.util;

import org.apache.avro.generic.GenericRecord;
import org.apache.avro.util.Utf8;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.parquet.avro.AvroParquetReader;
import org.apache.parquet.hadoop.ParquetReader;

import javax.sound.sampled.AudioFileFormat.Type;
import javax.sound.sampled.AudioFormat;
import javax.sound.sampled.AudioInputStream;
import javax.sound.sampled.AudioSystem;
import java.io.*;
import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.Paths;

public class ParquetProcessorTwo {

    private static String OUTPUT_DIR = "D:\\Darija-stt-mix\\train-00000-of-00001-b55cab9e7cf9a54b_2";

    public static void main(String[] args) {
        OUTPUT_DIR = "D:\\Darija-stt-mix\\train-00000-of-00001-b55cab9e7cf9a54b_2";
        String parquetFilePath = "D:\\Darija-stt-mix\\train-00000-of-00001-b55cab9e7cf9a54b_2.parquet";  // 请确保路径正确
        createOutputDir();
        readParquetFile(parquetFilePath);
    }

    private static void createOutputDir() {
        File dir = new File(OUTPUT_DIR);
        if (!dir.exists()) {
            if (dir.mkdirs()) {
                System.out.println("Output directory created at: " + OUTPUT_DIR);
            } else {
                System.err.println("Failed to create output directory at: " + OUTPUT_DIR);
            }
        }
    }

    private static void saveAudioFile(byte[] audioData, String fileName) throws IOException {
        System.out.println("Saving audio file to: " + fileName); // 调试信息

        if (audioData.length < 44) {
            throw new IOException("Invalid WAV file header, too short.");
        }

        ByteArrayInputStream bais = new ByteArrayInputStream(audioData);
        byte[] header = new byte[44];
        if (bais.read(header) != 44) {
            throw new IOException("Invalid WAV file header, unable to read 44 bytes.");
        }

        int sampleRate = ((header[24] & 0xFF) | ((header[25] & 0xFF) << 8) | ((header[26] & 0xFF) << 16) | ((header[27] & 0xFF) << 24));
        int bitsPerSample = ((header[34] & 0xFF) | ((header[35] & 0xFF) << 8));
        int channels = ((header[22] & 0xFF) | ((header[23] & 0xFF) << 8));

        AudioFormat format = new AudioFormat(sampleRate, bitsPerSample, channels, true, false);
        AudioInputStream audioInputStream = new AudioInputStream(bais, format, (audioData.length - 44) / format.getFrameSize());

        File audioFile = new File(fileName);
        audioFile.getParentFile().mkdirs(); // 确保目录存在
        AudioSystem.write(audioInputStream, Type.WAVE, audioFile);
    }

    private static void saveTextFile(String textData, String fileName) throws IOException {
        System.out.println("Saving text file to: " + fileName); // 调试信息
        File textFile = new File(fileName);
        textFile.getParentFile().mkdirs(); // 确保目录存在
        Files.write(Paths.get(fileName), textData.getBytes());
    }

    private static void logError(String errorMessage) {
        try (PrintWriter out = new PrintWriter(new BufferedWriter(new FileWriter("error.log", true)))) {
            out.println(errorMessage);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static void readParquetFile(String filePath) {
        File parquetFile = new File(filePath);
        if (!parquetFile.exists()) {
            System.err.println("File " + filePath + " does not exist.");
            return;
        } else {
            System.out.println("File " + filePath + " found.");
        }

        Path path = new Path(filePath);
        Configuration conf = new Configuration();

        try (ParquetReader<GenericRecord> reader = AvroParquetReader.<GenericRecord>builder(path).withConf(conf).build()) {

            GenericRecord record;
            int counter = 0;
            while ((record = reader.read()) != null) {
                try {
                    GenericRecord audioRecord = (GenericRecord) record.get("audio");
                    if (audioRecord != null) {
                        ByteBuffer audioByteBuffer = (ByteBuffer) audioRecord.get("bytes");
                        byte[] audioData = null;
                        if (audioByteBuffer != null) {
                            audioData = new byte[audioByteBuffer.remaining()];
                            audioByteBuffer.get(audioData);
                        }

                        String baseFileName = OUTPUT_DIR + File.separator + "output_" + counter;
                        String audioPath = baseFileName + System.currentTimeMillis() + ".wav";
                        String textPath = baseFileName+ System.currentTimeMillis() + ".txt";

                        if (audioData != null) {
                            saveAudioFile(audioData, audioPath);
                        }
                        saveTextFile( record.get("label")+"", textPath);
                        counter++;
                    }

                } catch (Exception e) {
                    logError("Error processing record #" + counter + ": " + e.getMessage());
                    e.printStackTrace(); // 打印错误堆栈信息
                }
            }

        } catch (IOException e) {
            logError("Error reading Parquet file: " + e.getMessage());
            e.printStackTrace(); // 打印错误堆栈信息
        }
    }
}
